#pragma once

#include <cmath>

#include "../common/HexagonalGrid.h"
#include "../common/GenericCA.h"

typedef float LBfloat;
const LBfloat eps = (LBfloat) 0.0001;

const LBfloat CellHeight = (LBfloat) (sqrt(3.0) / 2.0);

struct D2Q7Cell
{
  LBfloat CR, UR, UL, CL, DL, DR, CC;
};

inline LBfloat calcDensity(const D2Q7Cell & cell) 
{ 
  return cell.CR + cell.UR + cell.UL + cell.CL + cell.DL + cell.DR + cell.CC;
}

inline void calcMomentum( const D2Q7Cell & cell, LBfloat & ux, LBfloat & uy ) 
{ 
  ux = cell.CR - cell.CL + (LBfloat)0.5 * (cell.UR + cell.DR - cell.UL - cell.DL);
  uy = CellHeight * (cell.UR + cell.UL - cell.DR - cell.DL);
}

inline void calcVelocity( const D2Q7Cell & cell, LBfloat & ux, LBfloat & uy )
{
  calcMomentum(cell, ux, uy);
  LBfloat p = calcDensity(cell);
  if (p > eps)
  {
    ux /= p;
    uy /= p;
  }
}

enum CellType { CELL_EMPTY, CELL_WALL };

class LatticeBoltzmannD2Q7 : 
  public Generic2DCA<LatticeBoltzmannD2Q7, D2Q7Cell, CAHexagonalGrid>
{
  typedef Generic2DCA<LatticeBoltzmannD2Q7, D2Q7Cell, CAHexagonalGrid> Base;
public:
  LatticeBoltzmannD2Q7(int width, int height);

  void setConstBound(CellType cellType, const D2Q7Cell & cell);
  void setTorusBound();

  CellType getCellType(int x, int y) const { return m_cellTypes.getCell(x, y); }
  void setCellType(int x, int y, CellType cellType) { m_cellTypes.setCell(x, y, cellType); }

  void setViscosity(LBfloat v)
  {
    // v = (t - 0.5) / 4
    m_relaxCoef = (LBfloat) (1.0 / (4.0 * v + 0.5));
  }

  void calcEquilibrium(LBfloat p, LBfloat ux, LBfloat uy, D2Q7Cell & cell) const;

private:
  struct Visitor : public GridType::NeighbourVisitor<Visitor>
  {
    LatticeBoltzmannD2Q7 & ca;
    Visitor(LatticeBoltzmannD2Q7 & parent) : ca(parent) {}
    void walkCell();
  };

  LBfloat getRelaxCoef() const { return m_relaxCoef; }
  CellType getCellType(CellId id) const { return m_cellTypes.getCellById(id); }

  CAHexagonalGrid<CellType> m_cellTypes;
  LBfloat m_relaxCoef;
};
